import torch
import models_xh as models
import numpy as np
from PIL import Image
import cv2
from data_loader_xhy import prepare_image_cv2
import os


resume = './xhrs80_y256_138.pt'
img_dir = r'H:\data\XH_0111\labelme\01_bingsi_yin'
# img_dir = r'E:\xianhuang0818\xhsi_segall\img_yin'
w_dir = r'H:\data\XH_0111\seg_result'
# img_dir = r'G:\dataset\xh_seg_yin\images'
result_path = './examples/0.png'

model = models.resnet80(pretrained=False).cuda()
model.eval()

#resume..
checkpoint = torch.load(resume)
model.load_state_dict(checkpoint)

names = os.listdir(img_dir)
for name in names:
    title = name.split(".")[0]
    img_path = os.path.join(img_dir, name)
    img_u8 = cv2.imread(img_path, 1)
    cv2.imshow("img_u8", cv2.resize(img_u8, (600, 600)))
    img_u8 = cv2.cvtColor(img_u8, cv2.COLOR_BGR2GRAY)
    equ = cv2.equalizeHist(img_u8)
    img_u8 = cv2.cvtColor(equ, cv2.COLOR_GRAY2BGR)

    original_img = np.array(img_u8, dtype=np.float32)

    h, w, _ = original_img.shape

    img = prepare_image_cv2(original_img)
    img = torch.from_numpy(img).unsqueeze(0).cuda()

    outs = model(img)
    # result = outs[-1].squeeze().detach().cpu().numpy()
    result = outs.squeeze().detach().cpu().numpy()

    result = (result * 255).astype(np.uint8)
    cv2.imshow("result",cv2.resize(result, (600, 600)))

    # kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (40, 40))  # 矩形结构
    # result = cv2.morphologyEx(result, cv2.MORPH_OPEN, kernel)
    # cv2.imshow("open",cv2.resize(result, (600, 600)))
    # result = cv2.morphologyEx(result, cv2.MORPH_CLOSE, kernel)
    # cv2.imshow("close", cv2.resize(result, (600, 600)))

    imgShow = np.empty((600, 600, 3), np.uint8)
    equ_rs = cv2.resize(equ, (600, 600))
    imgShow[:, :, 0] = equ_rs
    imgShow[:, :, 1] = equ_rs
    imgShow[:, :, 2] = np.maximum(equ_rs, cv2.resize(result, (600, 600))) * 0.8 + equ_rs * 0.2
    cv2.imshow("imgShow", cv2.resize(imgShow, (600, 600)))
    cv2.waitKey(1000)
    # w_path = os.path.join(w_dir, "%s_0.jpg"%title)
    # cv2.imwrite(w_path, imgShow)
    # w_path = os.path.join(w_dir, "%s_1.jpg"%title)
    # cv2.imwrite(w_path, result)

    # Image.fromarray(result).save(result_path)

